// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/pipeline/py_torch_infer_op.h>

#include <tbir/pipeline/tx_session.h>

#include "userdata_mutator.h"

namespace tbir::runtime {

    TBIR_REGISTER_NATIVE_OP(PyTorchInferOp).SetThreadSafety(false);

    void PyTorchInferOp::Init() {
        auto impl_data = GetAttr<RTValue>("impl");
        UserDataMutator::Mutate(&impl_data, this);
        auto ud_ref = impl_data.AsObjectRef<UserDataRef>();
        impl_ = check_get_op_kernel(ud_ref);
        sub_ops_ = {impl_};
    }

    RTValue PyTorchInferOp::Process(PyArgs inputs) const {
        return impl_->Process(inputs);
    }

}  // namespace tbir::runtime
